{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "71eadd60",
   "metadata": {},
   "source": [
    "# Trajectory Follower Kinematic Simulation Sandbox\n",
    "\n",
    "Self-contained example of `home_robot.control.traj_following_controller.TrajFollower` in action.\n",
    "\n",
    "- The controller accepts a \"trajectory\" in the form of a function: `xyt, dxyt, done = traj(t)`.\n",
    "- The controller stabilizes around the trajectory, instead of moving directly to the goal location like the GotoVelocityController.\n",
    "- The trajectory can be updated mid-execution.\n",
    "\n",
    "### Requirements\n",
    "Pip\n",
    "- numpy\n",
    "- scipy\n",
    "- matplotlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2d69cb8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.animation as animation\n",
    "\n",
    "from home_robot.control.traj_following_controller import TrajFollower\n",
    "\n",
    "plt.rcParams[\"animation.html\"] = \"jshtml\"\n",
    "plt.rcParams['figure.dpi'] = 150"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ec1b8a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ----- HYPERPARAMS\n",
    "# Simulation params\n",
    "SPEEDUP = 4\n",
    "\n",
    "TIME_HORIZON = 12\n",
    "SIM_HZ = 60\n",
    "CONTROL_HZ = 20\n",
    "VIZ_HZ = 24 / SPEEDUP\n",
    "\n",
    "# Visualization params\n",
    "R_ROBOT = 0.2\n",
    "BUTT_SAMPLES = 20\n",
    "X_LIM = [-2, 2]\n",
    "Y_LIM = [-2, 2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b44df483",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ----- SIMULATION ENV\n",
    "class Env:\n",
    "    \"\"\"Simple kinematic sim of a diff drive robot\"\"\"\n",
    "    \n",
    "    def __init__(self, hz, sys_noise=None):\n",
    "        self.pos_state = np.zeros(3)  # x, y, t\n",
    "        self.vel_state = np.zeros(2)  # v, w\n",
    "        self.dt = 1.0 / hz\n",
    "        \n",
    "        self.noise = np.zeros(3) if sys_noise is None else np.array(sys_noise)\n",
    "        \n",
    "    def get_pose(self):\n",
    "        return self.pos_state\n",
    "    \n",
    "    def step(self, vel_input=None):\n",
    "        if vel_input is not None:\n",
    "            self.vel_state = vel_input\n",
    "            \n",
    "        self.pos_state[0] += self.vel_state[0] * np.cos(self.pos_state[2]) * self.dt\n",
    "        self.pos_state[1] += self.vel_state[0] * np.sin(self.pos_state[2]) * self.dt\n",
    "        self.pos_state[2] += self.vel_state[1] * self.dt\n",
    "        self.pos_state += self.noise * np.random.randn(3) * self.dt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20e86366",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ----- VISUALIZER\n",
    "class Viz:\n",
    "    \"\"\"Visualize and animate robot states\"\"\"\n",
    "\n",
    "    def __init__(self, title=\"\"):\n",
    "        self.title = title\n",
    "        self.frames = []\n",
    "        \n",
    "        # Robot local vertices\n",
    "        xy_ul = np.array([R_ROBOT / 2, R_ROBOT])\n",
    "        xy_ur = np.array([R_ROBOT / 2, -R_ROBOT])\n",
    "        xy_l = np.array([-R_ROBOT / 2, R_ROBOT])\n",
    "        xy_r = np.array([-R_ROBOT / 2, -R_ROBOT])\n",
    "        \n",
    "        xy_list = [xy_r, xy_ur, xy_ul, xy_l]\n",
    "        for phi in np.linspace(np.pi / 2, 3 * np.pi / 2, BUTT_SAMPLES):\n",
    "            xy_list.append(np.array([R_ROBOT * (np.cos(phi) - 0.5), R_ROBOT * np.sin(phi)]))\n",
    "            \n",
    "        self.xy_verts_local = np.array(xy_list)\n",
    "        \n",
    "        # Trajectory\n",
    "        traj_path = None\n",
    "        \n",
    "    def _draw_robot(self, xyt, color):\n",
    "        theta = xyt[2]\n",
    "        ct = np.cos(theta)\n",
    "        st = np.sin(theta)\n",
    "        trans = np.array([[ct, -st], [st, ct]])\n",
    "        \n",
    "        xy_verts_transformed = (trans @ self.xy_verts_local.T).T + xyt[None, :2]\n",
    "        \n",
    "        self.ax.plot(xyt[0], xyt[1], color=color, marker='o')\n",
    "        self.ax.plot(xy_verts_transformed[:, 0], xy_verts_transformed[:, 1], color=color)\n",
    "        \n",
    "    def _draw_traj(self, traj_path, color):\n",
    "        self.ax.plot(traj_path[:, 0], traj_path[:, 1], color=color)\n",
    "        \n",
    "    def _render_frame(self, frame):\n",
    "        plt.cla()\n",
    "        \n",
    "        self.ax.set_xlim(*X_LIM)\n",
    "        self.ax.set_ylim(*Y_LIM)\n",
    "        self.ax.set_aspect(\"equal\")\n",
    "        \n",
    "        if self.traj_path is not None:\n",
    "            self._draw_traj(self.traj_path, '0.5')\n",
    "        self._draw_robot(frame[1], '0.7')\n",
    "        self._draw_robot(frame[0], '0')\n",
    "        \n",
    "    def clear_frames(self):\n",
    "        self.frames = []\n",
    "        \n",
    "    def set_trajectory(self, traj_func, T, N):\n",
    "        \"\"\"Assign trajectory at beginning of episode to visualize desired states\"\"\"\n",
    "        self.traj_path = np.zeros([N, 2])\n",
    "        for i in range(N):\n",
    "            t = i * T / N\n",
    "            xyt, _, _ = traj_func(t)\n",
    "            self.traj_path[i, :] = xyt[:2]\n",
    "    \n",
    "    def record_frame(self, xyt_robot, xyt_goal):\n",
    "        assert xyt_robot.shape[-1] == 3\n",
    "        assert xyt_goal.shape[-1] == 3\n",
    "        self.frames.append((xyt_robot, xyt_goal))\n",
    "        \n",
    "    def animate(self, playback_speed=1.0):\n",
    "        \"\"\"Animate all recorded frames\"\"\"\n",
    "        plt.ioff()\n",
    "        fig, self.ax = plt.subplots()\n",
    "        self.ax.set_title(self.title)\n",
    "        \n",
    "        return animation.FuncAnimation(fig, self._render_frame, frames=self.frames, interval=1000 / VIZ_HZ / playback_speed, repeat=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6b21681",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ----- GENERATE TRAJ: Circle\n",
    "R = 1.\n",
    "\n",
    "w = np.pi / 2 / TIME_HORIZON\n",
    "\n",
    "def traj_func(t):\n",
    "    \"\"\"A trajectory function that returns the desired SE2 position & velocity given a timestamp\"\"\"\n",
    "    t_scaled = t / TIME_HORIZON\n",
    "    theta = t_scaled * (np.pi / 2)\n",
    "    \n",
    "    xyt =  np.array([\n",
    "        R * np.sin(theta),\n",
    "        R * (1. - np.cos(theta)),\n",
    "        theta\n",
    "    ])\n",
    "    dxyt = np.array([\n",
    "        R * w * np.cos(theta),\n",
    "        R * w * np.sin(theta),\n",
    "        w\n",
    "    ])\n",
    "    done = t >= TIME_HORIZON\n",
    "    \n",
    "    return xyt, dxyt, done"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb3dc36c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ----- CONTROLLER\n",
    "from omegaconf import OmegaConf\n",
    "\n",
    "def create_controller(\n",
    "    v_max=0.2,\n",
    "    w_max=0.45,\n",
    "    k_p=0.5,\n",
    "    damp_ratio=1.0,\n",
    "    decay=0.5,\n",
    "):\n",
    "    # Generate config\n",
    "    cfg_dict = {\n",
    "        \"v_max\": v_max,\n",
    "        \"w_max\": w_max,\n",
    "        \"k_p\": k_p,\n",
    "        \"damp_ratio\": damp_ratio,\n",
    "        \"decay\": decay,\n",
    "    }\n",
    "    cfg = OmegaConf.create(cfg_dict)\n",
    "\n",
    "    # Instatiate controller\n",
    "    return TrajFollower(cfg)  # default config \"traj_follower.yaml\" will be loaded if no cfg is provided"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c25b035c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ----- MAIN\n",
    "OPEN_LOOP = False  # change this value to compare using the TrajFollower vs. open-loop execution of the trajectory\n",
    "\n",
    "# initialize objects\n",
    "xyt_goal, _, _ = traj_func(TIME_HORIZON)\n",
    "\n",
    "viz = Viz()\n",
    "env = Env(SIM_HZ, sys_noise=[0.03, 0.03, 0.2])  # inject noise into system to observe closed-loop control behavior\n",
    "controller = create_controller()\n",
    "\n",
    "# set traj\n",
    "controller.update_trajectory(traj_func)  # Assign trajectory to controller\n",
    "viz.set_trajectory(traj_func, TIME_HORIZON, int(TIME_HORIZON * SIM_HZ))\n",
    "\n",
    "# simulate\n",
    "t_control = 0\n",
    "t_viz = 0\n",
    "dt_control = 1 / CONTROL_HZ\n",
    "dt_viz = 1 / VIZ_HZ\n",
    "for t in np.linspace(0, TIME_HORIZON, int(TIME_HORIZON * SIM_HZ)):\n",
    "    xyt = env.get_pose()\n",
    "\n",
    "    vel_command = None\n",
    "    \n",
    "    # Query controller at a frequency of CONTROL_HZ\n",
    "    if t >= t_control:\n",
    "        # Open loop control: directly apply desired velocities\n",
    "        if OPEN_LOOP:\n",
    "            _, dxyt_goal, _ = traj_func(t)\n",
    "            vel_command = [np.linalg.norm(dxyt_goal[:2]), dxyt_goal[2]]\n",
    "    \n",
    "        # Closed loop control: Use traj follower\n",
    "        else:\n",
    "            vel_command = controller.forward(xyt, t)\n",
    "        t_control += dt_control\n",
    "        \n",
    "    # Record frames at a frequency of VIZ_HZ\n",
    "    if t >= t_viz:\n",
    "        xyt_goal, _, _ = traj_func(t)\n",
    "        viz.record_frame(xyt.copy(), xyt_goal.copy())\n",
    "        t_viz += dt_viz\n",
    "\n",
    "    env.step(vel_command)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "128d46ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "# animate\n",
    "viz.animate(playback_speed=SPEEDUP)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8a5a1ae",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
